import os
import re
import json
import time
import uuid
import base64
from decimal import Decimal
from datetime import datetime, date
from flask import Flask, request, jsonify, send_from_directory
from flask_sqlalchemy import SQLAlchemy
from werkzeug.security import generate_password_hash, check_password_hash
from werkzeug.utils import secure_filename
from flask_cors import CORS
from openai import OpenAI

app = Flask('发票系统')
CORS(app, resources={r"/*": {"origins": "*"}})

# 配置数据库
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://root:123456@localhost/invoice_management'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB

# 初始化OpenAI客户端
client = OpenAI(
    api_key="sk-52b4549679e24d04854f1465b9360837",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

db = SQLAlchemy(app)

# 创建模型
class User(db.Model):
    __tablename__ = 'users'
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    password = db.Column(db.String(255), nullable=False)
    created_at = db.Column(db.DateTime, server_default=db.func.now())

class Invoice(db.Model):
    __tablename__ = 'invoices'
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    invoice_number = db.Column(db.String(50), nullable=False)
    invoice_date = db.Column(db.Date, nullable=False)
    amount = db.Column(db.Numeric(10, 2), nullable=False)
    tax_amount = db.Column(db.Numeric(10, 2))
    total_amount = db.Column(db.Numeric(10, 2), nullable=False)
    seller_name = db.Column(db.String(100))
    seller_tax_id = db.Column(db.String(50))
    buyer_name = db.Column(db.String(100))
    buyer_tax_id = db.Column(db.String(50))
    invoice_type = db.Column(db.String(20))
    items = db.Column(db.Text)
    image_path = db.Column(db.String(255), nullable=False)
    created_at = db.Column(db.DateTime, server_default=db.func.now())

# 辅助函数
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in {'png', 'jpg', 'jpeg'}

def clean_amount(amount_str):
    """清理金额字符串，移除逗号并转换为Decimal"""
    if not amount_str:
        return Decimal('0.00')
    try:
        cleaned = re.sub(r'[^\d.-]', '', str(amount_str))
        return Decimal(cleaned).quantize(Decimal('0.00'))
    except:
        return Decimal('0.00')

def parse_date(date_str):
    """解析日期字符串，确保为YYYY-MM-DD格式"""
    try:
        if not date_str:
            return None
        return datetime.strptime(date_str, '%Y-%m-%d').date()
    except ValueError:
        return None

def extract_json_from_response(response_text):
    """从API响应中提取和修复JSON数据"""
    if not response_text or not isinstance(response_text, str):
        raise ValueError("无效的API响应")
    
    # 尝试直接解析
    try:
        return json.loads(response_text)
    except json.JSONDecodeError:
        pass
    
    # 如果直接解析失败，尝试提取JSON部分
    json_start = response_text.find('{')
    json_end = response_text.rfind('}') + 1
    if json_start == -1 or json_end == 0:
        raise ValueError("响应中未找到JSON数据")
    
    json_str = response_text[json_start:json_end]
    
    # 尝试修复常见的JSON格式问题
    repair_attempts = [
        # 尝试1: 替换单引号为双引号
        lambda s: s.replace("'", '"'),
        # 尝试2: 确保键名有引号
        lambda s: re.sub(r'(\w+)\s*:', r'"\1":', s),
        # 尝试3: 确保字符串值有引号
        lambda s: re.sub(r':\s*([^"\d][^,}]*)([,}])', lambda m: f': "{m.group(1)}"{m.group(2)}', s),
        # 尝试4: 移除注释
        lambda s: re.sub(r'//.*?\n', '', s),
        # 尝试5: 修复布尔值和null
        lambda s: s.replace(': true', ': true').replace(': false', ': false').replace(': null', ': null')
    ]
    
    for attempt in repair_attempts:
        try:
            repaired = attempt(json_str)
            return json.loads(repaired)
        except json.JSONDecodeError:
            continue
    
    # 最终尝试: 移除所有非JSON字符
    try:
        cleaned = re.sub(r'[^\x20-\x7E]', '', json_str)
        return json.loads(cleaned)
    except json.JSONDecodeError as e:
        raise ValueError(f"无法解析API响应: {str(e)}")

def recognize_invoice_with_qwen(image_path, max_retries=3):
    """识别发票信息，带有重试机制"""
    retry_count = 0
    last_error = None
    
    while retry_count < max_retries:
        try:
            with open(image_path, 'rb') as f:
                image_data = f.read()
            
            base64_image = base64.b64encode(image_data).decode('utf-8')
            
            prompt = '''请精确识别这张发票上的所有字段并以严格JSON格式返回：
            {
                "invoice_number": "发票号码",
                "invoice_date": "开票日期(YYYY-MM-DD格式)",
                "amount": "金额(小写数字)",
                "tax_amount": "税额(小写数字)",
                "total_amount": "价税合计金额(小写数字)",
                "seller_name": "销售方名称",
                "seller_tax_id": "销售方税号",
                "buyer_name": "购买方名称",
                "buyer_tax_id": "购买方税号",
                "invoice_type": "发票类型",
                "items": [
                    {
                        "name": "商品名称",
                        "spec": "规格型号",
                        "quantity": "数量",
                        "price": "单价",
                        "amount": "金额",
                        "tax_rate": "税率"
                    }
                ]
            }'''
            
            completion = client.chat.completions.create(
                model="qwen-vl-plus",
                messages=[{
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
                    ]
                }],
                temperature=0.1,
                top_p=0.9
            )
            
            content = completion.choices[0].message.content
            
            # 记录原始响应
            app.logger.info(f"API原始响应: {content}")
            
            # 提取和验证JSON
            invoice_data = extract_json_from_response(content)
            
            # 验证必要字段
            required_fields = ['invoice_number', 'invoice_date', 'total_amount']
            for field in required_fields:
                if field not in invoice_data:
                    raise ValueError(f"缺少必要字段: {field}")
            
            # 处理金额字段
            def process_amount(amount_str, field_name):
                try:
                    if isinstance(amount_str, (int, float, Decimal)):
                        return Decimal(str(amount_str)).quantize(Decimal('0.00'))
                    return clean_amount(str(amount_str))
                except:
                    app.logger.warning(f"无法解析{field_name}值: {amount_str}")
                    return Decimal('0.00')
            
            # 处理主要金额
            invoice_data['amount'] = process_amount(invoice_data.get('amount', 0), 'amount')
            invoice_data['tax_amount'] = process_amount(invoice_data.get('tax_amount', 0), 'tax_amount')
            invoice_data['total_amount'] = process_amount(invoice_data.get('total_amount', 0), 'total_amount')
            
            # 计算缺失的金额
            if invoice_data['tax_amount'] == Decimal('0.00') and invoice_data['total_amount'] > invoice_data['amount']:
                invoice_data['tax_amount'] = invoice_data['total_amount'] - invoice_data['amount']
            elif invoice_data['total_amount'] == Decimal('0.00') and invoice_data['tax_amount'] > Decimal('0.00'):
                invoice_data['total_amount'] = invoice_data['amount'] + invoice_data['tax_amount']
            
            # 处理商品项
            if 'items' in invoice_data:
                for item in invoice_data['items']:
                    item['price'] = str(process_amount(item.get('price', 0), 'price'))
                    item['amount'] = str(process_amount(item.get('amount', 0), 'amount'))
                    if 'tax_rate' in item:
                        item['tax_rate'] = str(re.sub(r'[^\d.]', '', str(item['tax_rate'])))
            
            return invoice_data
            
        except Exception as e:
            last_error = e
            retry_count += 1
            time.sleep(1)
            app.logger.error(f"识别尝试 {retry_count}/{max_retries} 失败: {str(e)}", exc_info=True)
    
    raise ValueError(f"经过{max_retries}次尝试后仍失败: {str(last_error)}")

# API路由
@app.route('/api/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    
    if not username or not password:
        return jsonify({'error': '用户名和密码不能为空'}), 400
    
    if User.query.filter_by(username=username).first():
        return jsonify({'error': '用户名已存在'}), 400
    
    hashed_password = generate_password_hash(password)
    new_user = User(username=username, password=hashed_password)
    db.session.add(new_user)
    db.session.commit()
    
    return jsonify({'message': '注册成功'}), 201

@app.route('/api/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    
    user = User.query.filter_by(username=username).first()
    
    if not user or not check_password_hash(user.password, password):
        return jsonify({'error': '用户名或密码错误'}), 401
    
    return jsonify({'message': '登录成功', 'user_id': user.id}), 200

@app.route('/api/upload', methods=['POST'])
def upload_invoice():
    if 'file' not in request.files:
        return jsonify({'error': '没有文件上传'}), 400
    
    file = request.files['file']
    user_id = request.form.get('user_id')
    
    if not user_id:
        return jsonify({'error': '用户ID不能为空'}), 400
    
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    if file and allowed_file(file.filename):
        if not os.path.exists(app.config['UPLOAD_FOLDER']):
            os.makedirs(app.config['UPLOAD_FOLDER'])
        
        filename = secure_filename(file.filename)
        unique_filename = f"{uuid.uuid4().hex}_{filename}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
        file.save(filepath)
        
        try:
            invoice_info = recognize_invoice_with_qwen(filepath)
            
            return jsonify({
                'message': '发票上传成功',
                'image_path': unique_filename,
                'invoice_info': invoice_info
            }), 200
        except Exception as e:
            if os.path.exists(filepath):
                os.remove(filepath)
            app.logger.error(f"发票识别失败: {str(e)}", exc_info=True)
            return jsonify({
                'error': '发票识别失败',
                'details': str(e),
                'advice': '请确保上传的是清晰的发票图片，或稍后再试'
            }), 500
    else:
        return jsonify({'error': '不支持的文件类型'}), 400

@app.route('/api/save_invoice', methods=['POST'])
def save_invoice():
    data = request.get_json()
    user_id = data.get('user_id')
    invoice_info = data.get('invoice_info')
    image_path = data.get('image_path')
    
    if not user_id or not invoice_info or not image_path:
        return jsonify({'error': '缺少必要参数'}), 400
    
    try:
        if not invoice_info.get('invoice_number'):
            return jsonify({'error': '发票号码不能为空'}), 400
        
        invoice_date = parse_date(invoice_info['invoice_date'])
        if not invoice_date:
            return jsonify({'error': '开票日期格式错误，请使用YYYY-MM-DD格式'}), 400
        
        amount = clean_amount(invoice_info.get('amount', '0'))
        tax_amount = clean_amount(invoice_info.get('tax_amount', '0'))
        total_amount = clean_amount(invoice_info.get('total_amount', '0'))
        
        if total_amount == Decimal('0.00'):
            total_amount = amount + tax_amount
        elif tax_amount == Decimal('0.00') and total_amount > amount:
            tax_amount = total_amount - amount
        
        new_invoice = Invoice(
            user_id=user_id,
            invoice_number=invoice_info.get('invoice_number', ''),
            invoice_date=invoice_date,
            amount=amount,
            tax_amount=tax_amount,
            total_amount=total_amount,
            seller_name=invoice_info.get('seller_name', ''),
            seller_tax_id=invoice_info.get('seller_tax_id', ''),
            buyer_name=invoice_info.get('buyer_name', ''),
            buyer_tax_id=invoice_info.get('buyer_tax_id', ''),
            invoice_type=invoice_info.get('invoice_type', ''),
            items=json.dumps(invoice_info.get('items', [])),
            image_path=image_path
        )
        
        db.session.add(new_invoice)
        db.session.commit()
        
        return jsonify({'message': '发票保存成功', 'invoice_id': new_invoice.id}), 201
    except Exception as e:
        db.session.rollback()
        app.logger.error(f"保存发票失败: {str(e)}", exc_info=True)
        return jsonify({'error': f'保存发票失败: {str(e)}'}), 500

@app.route('/api/query', methods=['GET'])
def query_invoices():
    user_id = request.args.get('user_id')
    start_date = request.args.get('start_date')
    end_date = request.args.get('end_date')
    keyword = request.args.get('keyword')
    
    if not user_id:
        return jsonify({'error': '用户ID不能为空'}), 400
    
    query = Invoice.query.filter_by(user_id=user_id)
    
    if start_date and end_date:
        try:
            start = datetime.strptime(start_date, '%Y-%m-%d').date()
            end = datetime.strptime(end_date, '%Y-%m-%d').date()
            query = query.filter(Invoice.invoice_date.between(start, end))
        except:
            return jsonify({'error': '日期格式错误，请使用YYYY-MM-DD格式'}), 400
    
    if keyword:
        keyword = f"%{keyword}%"
        query = query.filter(
            db.or_(
                Invoice.invoice_number.like(keyword),
                Invoice.seller_name.like(keyword),
                Invoice.buyer_name.like(keyword),
                Invoice.invoice_type.like(keyword)
            )
        )
    
    invoices = query.order_by(Invoice.invoice_date.desc()).all()
    
    result = []
    for invoice in invoices:
        result.append({
            'id': invoice.id,
            'invoice_number': invoice.invoice_number,
            'invoice_date': invoice.invoice_date.strftime('%Y-%m-%d') if invoice.invoice_date else None,
            'amount': float(invoice.amount) if invoice.amount else None,
            'total_amount': float(invoice.total_amount) if invoice.total_amount else None,
            'seller_name': invoice.seller_name,
            'buyer_name': invoice.buyer_name,
            'invoice_type': invoice.invoice_type,
            'image_path': invoice.image_path,
            'items': json.loads(invoice.items) if invoice.items else []
        })
    
    return jsonify({'invoices': result}), 200

@app.route('/api/statistics', methods=['GET'])
def get_statistics():
    user_id = request.args.get('user_id')
    start_date = request.args.get('start_date')
    end_date = request.args.get('end_date')
    keyword = request.args.get('keyword')
    
    if not user_id:
        return jsonify({'error': '用户ID不能为空'}), 400
    
    query = Invoice.query.filter_by(user_id=user_id)
    
    if start_date and end_date:
        try:
            start = datetime.strptime(start_date, '%Y-%m-%d').date()
            end = datetime.strptime(end_date, '%Y-%m-%d').date()
            query = query.filter(Invoice.invoice_date.between(start, end))
        except:
            return jsonify({'error': '日期格式错误，请使用YYYY-MM-DD格式'}), 400
    
    if keyword:
        keyword = f"%{keyword}%"
        query = query.filter(
            db.or_(
                Invoice.invoice_number.like(keyword),
                Invoice.seller_name.like(keyword),
                Invoice.buyer_name.like(keyword),
                Invoice.invoice_type.like(keyword)
            )
        )
    
    count = query.count()
    total_amount = query.with_entities(db.func.sum(Invoice.total_amount)).scalar() or 0
    
    return jsonify({
        'count': count,
        'total_amount': float(total_amount)
    }), 200

@app.route('/uploads/<filename>')
def get_invoice_image(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

if __name__ == '__main__':
    with app.app_context():
        db.create_all()
    app.run(debug=True, host='127.0.0.1', port=8080)
